import base64
import json
import re
from urllib.parse import parse_qs, urlparse

from bs4 import BeautifulSoup
from curl_cffi import requests
from flask import abort

from app.chatgpt.arkose import Arkose
from app.chatgpt.proofowork import ChatGPTPoW
from app.chatgpt.utils import generate_code_challenge, generate_code_verifier
from app.exceptions import ApiException
from app.utils.decorator import cache_func_data

user_agent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36"


class AuthWebIos:
    def __init__(self):
        self.session = requests.Session()
        self.session.headers.update({
            "content-type": "application/x-www-form-urlencoded",
            "User-Agent": user_agent,
            "accept-language": "zh-CN,zh;q=0.9"
        })
        self.session.impersonate = 'safari'
        self.session.timeout = 10

    def set_arkose_token(self, blob_value):
        arkose_json = Arkose(user_agent).get_arkose_token(blob_value)

        if "sup=1" not in arkose_json["token"]:
            raise ApiException("login web", "获取 arkoseToken 出错，待接入第三方打码平台")
        self.session.cookies.set('arkoseToken', arkose_json["token"])

    @cache_func_data(cache_time=900)
    def get_preauth_cookie(self):
        url = "https://publicapi.closeai.biz/auth/preauth"
        res = requests.get(url).json()
        return res["preauth_cookie"]

    def _login(self, username, password):
        code_verifier = generate_code_verifier()
        code_challenge = generate_code_challenge(code_verifier)

        base_url = "https://auth0.openai.com/authorize?"
        client_info = {"version": "2.9.0-beta.0", "name": "Auth0.swift", "env": {"swift": "5.x", "macOS": "14.4"}}
        params = {
            "response_type": "code",
            # "state": generate_random_string(43),
            # "ios_device_id": "62BFAD6A-E648-4675-80AD-300268AD6010",
            "client_id": "pdlLIX2Y72MIl2rhLhTE9VV9bN905kBh",
            "code_challenge": code_challenge,
            "code_challenge_method": "S256",
            "scope": "openid%20email%20profile%20offline_access%20model.request%20model.read%20organization.read%20organization.write",
            "redirect_uri": "com.openai.chat://auth0.openai.com/ios/com.openai.chat/callback",
            "audience": "https://api.openai.com/v1",
            # "preauth_cookie": "62BFAD6A-E648-4675-80AD-300268AD6010:1724046725-VbsBdMxucRpMCKkLFZJbXaq5DJeADdMZsVco2A%252Bnj68%253D",
            "preauth_cookie": self.get_preauth_cookie(),
            "ios_app_version": "1723258889",
            "prompt": "login",
            "auth0Client": base64.b64encode(json.dumps(client_info).replace(" ", "").encode()).decode()
        }

        query_string = "&".join([f"{key}={value}" for key, value in params.items()])
        login_url = base_url + query_string
        # print(login_url)

        identifier_res = self.session.get(login_url)
        self.session.headers.update({'referer': identifier_res.url})
        username_page_url = identifier_res.url

        # 输入账号
        identifier_query_params = parse_qs(urlparse(username_page_url).query)
        password_res = self.session.post(username_page_url, data={
            "state": identifier_query_params['state'][0],
            "username": username,
            "js-available": True,
            "webauthn-available": True,
            "is-brave": False,
            "webauthn-platform-available": True,
            "action": "default"
        })

        if password_res.status_code == 400:
            soup = BeautifulSoup(password_res.text, 'html.parser')
            err_span = soup.find('span', {'id': 'error-element-username'})
            abort(400, err_span.get_text().strip())

        # 获取 auth0-client-id
        pattern = re.compile(r'document\.cookie\s*=\s*`[^;]*auth0-client-id=([^;]+);')
        match = pattern.search(password_res.text)
        auth0_client_id = match.group(1)
        self.session.cookies.set('auth0-client-id', auth0_client_id)
        # self.session.cookies.set('oai-did', device_id)

        # 获取 data[blob] 并获取 arkose token
        pattern = re.compile(r'blob:\s*"([^"]+)"')
        match = pattern.search(password_res.text)
        blob_value = match.group(1)
        self.set_arkose_token(blob_value)
        # set_sentinel_token(device_id)

        # 输入密码
        password_query_params = parse_qs(urlparse(password_res.url).query)
        self.session.headers.update({'referer': password_res.url})
        home1_res = self.session.post(password_res.url, data={
            "state": password_query_params['state'][0],
            "username": username,
            "password": password
        }, allow_redirects=False)

        if home1_res.status_code == 400:
            soup = BeautifulSoup(home1_res.text, 'html.parser')
            err_span = soup.find('span', {'id': 'error-element-password'})
            abort(400, err_span.get_text().strip())

        location1 = home1_res.headers['Location']
        # print("location1", location1)
        home2_res = self.session.get('https://auth0.openai.com' + location1, allow_redirects=False)

        location2 = home2_res.headers['Location']
        # print("location2", location2)
        url_params = parse_qs(urlparse(location2).query)
        url = 'https://auth0.openai.com/oauth/token'

        data = {
            'redirect_uri': 'com.openai.chat://auth0.openai.com/ios/com.openai.chat/callback',
            'grant_type': 'authorization_code',
            'client_id': 'pdlLIX2Y72MIl2rhLhTE9VV9bN905kBh',
            'code': url_params['code'][0],
            'code_verifier': code_verifier,
        }
        resp = self.session.post(url, json=data, allow_redirects=False)
        return resp.json()

    @cache_func_data(cache_time=600)
    # @cache_func_data(cache_time=86400)
    def get_token_info(self, username, password):
        return self._login(username, password)

    def set_sentinel_token(self, device_id):
        chatgpt_pow = ChatGPTPoW()

        sentinel_res = self.session.post(
            "https://chatgpt.com/backend-api/sentinel/req",
            json={"id": device_id, "p": "gAAAAAC", "flow": "login"}
        )
        sentinel_json = sentinel_res.json()
        p = chatgpt_pow.get_answer({
            'seed': sentinel_json["proofofwork"]["seed"],
            'difficulty': sentinel_json["proofofwork"]["difficulty"],
            'user_agent': user_agent
        })
        sentinelToken = {"c": sentinel_json['token'], "p": p}
        self.session.cookies.set('sentinelToken', json.dumps(sentinelToken))

    def get_login_url(self, device_id, csrf_token):
        login_web_url = "https://chatgpt.com/api/auth/signin/login-web"
        params = {
            "prompt": "login",
            "screen_hint": "login",
            "ext-statsig-tier": "production",
            "ext-oai-did": device_id,
            "flow": "control"
        }
        data = {
            "callbackUrl": "/",
            "csrfToken": csrf_token.split("%7C")[0],
            "json": "true"
        }
        login_web_res = self.session.post(login_web_url, params=params, data=data)
        url = login_web_res.json()["url"]
        return url.replace("https://auth.openai.com", "https://auth0.openai.com")

    def __repr__(self):
        return self.__class__.__name__
